import random

import cv2
import os
import numpy as np
import albumentations as A

def load_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

def load_bboxes(bbox_path, image_shape):
    bboxes = []
    height, width, _ = image_shape
    with open(bbox_path, 'r') as file:
        lines = file.readlines()
        for line in lines:
            label, x_center, y_center, bbox_width, bbox_height = map(float, line.strip().split())
            x_min = (x_center - bbox_width / 2) * width
            y_min = (y_center - bbox_height / 2) * height
            x_max = (x_center + bbox_width / 2) * width
            y_max = (y_center + bbox_height / 2) * height
            bboxes.append([x_min, y_min, x_max, y_max, label])
    return bboxes


def save_image(image, output_folder, filename):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    image_path = os.path.join(output_folder, filename)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(image_path, image)

def save_bboxes(bboxes, output_folder, filename, image_shape):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    bbox_path = os.path.join(output_folder, filename)
    height, width, _ = image_shape
    with open(bbox_path, 'w') as file:
        for bbox in bboxes:
            x_min, y_min, x_max, y_max, label = bbox
            x_center = ((x_min + x_max) / 2) / width
            y_center = ((y_min + y_max) / 2) / height
            bbox_width = (x_max - x_min) / width
            bbox_height = (y_max - y_min) / height
            file.write(f"{int(label)} {x_center} {y_center} {bbox_width} {bbox_height}\n")


def calculate_crop_area(image_shape, bboxes, margin=0.1):
    height, width, _ = image_shape

    if not bboxes:
        return 0, 0, width, height

    # 找到包含所有标签的最小矩形
    x_min = min(bbox[0] for bbox in bboxes)
    # print(x_min)
    y_min = min(bbox[1] for bbox in bboxes)
    x_max = max(bbox[2] for bbox in bboxes)
    y_max = max(bbox[3] for bbox in bboxes)

    # 添加边距
    x_min = max(0, x_min - width * margin)
    y_min = max(0, y_min - height * margin)
    x_max = min(width, x_max + width * margin)
    y_max = min(height, y_max + height * margin)

    return x_min, y_min, x_max, y_max

# 亮度变化后裁剪
# def augment_brightness_contrast(image, bboxes, image_shape):  # 亮度变化后裁剪
#     if not bboxes:
#         return image, bboxes
#
#     crop_area = calculate_crop_area(image_shape, bboxes)
#
#     # 将标签和边界框分开，因为 albumentations 需要它们分开
#     bboxes, labels = zip(*[(bbox[:-1], bbox[-1]) for bbox in bboxes])
#     transform = A.Compose([
#         A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=1),
#         A.Crop(*map(int, crop_area), p=1)
#     ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
#
#     augmented = transform(image=image, bboxes=bboxes, labels=labels)
#     # 将边界框和标签重新组合为一个列表
#     augmented_bboxes = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]
#     return augmented['image'], augmented_bboxes

def crop_image(image, bboxes, image_shape):
    if not bboxes:
        return image, bboxes
    # 将标签和边界框分开
    bboxes, labels = zip(*[(bbox[:-1], bbox[-1]) for bbox in bboxes])
    # bboxes:([528.99944, 653.000064, 867.0002400000001, 1067.000496],)
    img_height, img_width = image.shape[:2]

    def clip_bbox(bboxes, img_width, img_height):
        clipped_bboxes = []
        for bbox in bboxes:
            # """将边界框修剪到图像范围内。"""
            x_min, y_min, x_max, y_max = bbox
            x_min = max(0, min(x_min, img_width))
            y_min = max(0, min(y_min, img_height))
            x_max = max(0, min(x_max, img_width))
            y_max = max(0, min(y_max, img_height))
            clipped_bboxes.append((x_min, y_min, x_max, y_max))
            # print(f'clipped_bboxes:{clipped_bboxes}')
            # print('--------------------')
        return clipped_bboxes

    bboxes = clip_bbox(bboxes, img_width, img_height)
    crop_area = calculate_crop_area(image_shape, bboxes)
    transform = A.Compose([
        A.Crop(*map(int, crop_area), p=1)
    ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
    augmented = transform(image=image, bboxes=bboxes, labels=labels)
    # 重新组合边界框和标签
    augmented_bboxes = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]
    # augmented_bboxes:[(232.99944000000005, 171.00006399999995, 571.0002400000001, 585.0004960000001, 0.0)]
    # print(f'augmented_bboxes:{augmented_bboxes}')
    # print('--------------------')
    return augmented['image'], augmented_bboxes



#  裁剪调试失败
# def crop_image(image, bboxes_with_labels, image_shape):
#     if not bboxes_with_labels:
#         return image, bboxes_with_labels
#
#     bboxes, labels = zip(*bboxes_with_labels)
#     # bboxes:([528.99944, 653.000064, 867.0002400000001, 1067.000496],)
#     print(f'bboxes:{bboxes}')
#     print('--------------------')
#
#     # def clip_bboxes(bboxes):
#     #     clipped_bboxes = []
#     #     for bbox in bboxes:
#     #         x_min, y_min, x_max, y_max, label = bbox
#     #         x_min, y_min, x_max, y_max = np.clip([x_min, y_min, x_max, y_max], 0.0, 1.0)
#     #         clipped_bboxes.append((x_min, y_min, x_max, y_max, label))
#     #     return clipped_bboxes
#     #
#     # bboxes = clip_bboxes(bboxes)
#     # 计算裁剪区域
#     crop_area = calculate_crop_area(image_shape, bboxes)
#
#
#     transform = A.Compose([
#         A.Crop(*map(int, crop_area), p=1)
#     ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
#
#
#     augmented = transform(image=image, bboxes=bboxes, labels=labels)
#     # 重新组合边界框和标签
#     augmented_bboxes_with_labels = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]
#     return augmented['image'], augmented_bboxes_with_labels
#

def change_brightness_contrast(image, bboxes, brightness_value=0, contrast_value=1):
    bboxes_with_labels = [(bbox[:-1], bbox[-1]) for bbox in bboxes]
    if not bboxes_with_labels:
        return image, bboxes_with_labels

    augmented_bboxes = [(x_min, y_min, x_max, y_max, label) for ([x_min, y_min, x_max, y_max], label) in
                        bboxes_with_labels]
    # augmented_bboxes:[(528.99944, 653.000064, 867.0002400000001, 1067.000496, 0.0)]
    # print(f'augmented_bboxes:{augmented_bboxes}')
    # 将图像从[0, 255]范围转换到[0, 1]范围
    image = image.astype(np.float32) / 255.0
    # 调整亮度
    image = image + brightness_value
    # 调整对比度
    image = image * contrast_value
    # 将图像剪切回[0, 1]范围
    image = np.clip(image, 0, 1)
    # 将图像从[0, 1]范围转换回[0, 255]范围
    image = (image * 255).astype(np.uint8)
    return image, augmented_bboxes
#随机亮度
# def random_brightness(image, bboxes, value=0.2):
#     transform = A.Compose([
#         A.RandomBrightnessContrast(brightness_limit=value, contrast_limit=value, p=1)
#     ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
#
#     # 将标签和边界框分开
#     bboxes, labels = zip(*[(bbox[:-1], bbox[-1]) for bbox in bboxes])
#     augmented = transform(image=image, bboxes=bboxes, labels=labels)
#
#     # 将边界框和标签重新组合为一个列表
#     augmented_bboxes = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]
#     return augmented['image'], augmented_bboxes


def augment_noise(image, bboxes):
    if not bboxes:
        return image, bboxes
    # 将标签和边界框分开
    bboxes, labels = zip(*[(bbox[:-1], bbox[-1]) for bbox in bboxes])
    # bboxes:([528.99944, 653.000064, 867.0002400000001, 1067.000496],)
    img_height, img_width = image.shape[:2]

    def clip_bbox(bboxes, img_width, img_height):
        clipped_bboxes = []
        for bbox in bboxes:
        # """将边界框修剪到图像范围内。"""
            x_min, y_min, x_max, y_max = bbox
            x_min = max(0, min(x_min, img_width))
            y_min = max(0, min(y_min, img_height))
            x_max = max(0, min(x_max, img_width))
            y_max = max(0, min(y_max, img_height))
            clipped_bboxes.append((x_min, y_min, x_max, y_max))
            # print(f'clipped_bboxes:{clipped_bboxes}')
            # print('--------------------')
        return clipped_bboxes
    bboxes = clip_bbox(bboxes, img_width, img_height)
    # print(f'bboxes:{bboxes}')
    # print('--------------------')
    # print(f'函数里面bboxes:{bboxes}\n')
    # [center_x, center_y, w, h] = [*bboxes]
    transform = A.Compose([
        A.GaussNoise(p=1)
    ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))
    # print(f'bboxes:{bboxes}')
    # print('--------------------')
    augmented = transform(image=image, bboxes=bboxes, labels=labels)
    # 将边界框和标签重新组合为一个列表
    augmented_bboxes = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]
    # print(f'augmented_bboxes:{augmented_bboxes}')
    return augmented['image'], augmented_bboxes


def augment_rotation(image, bboxes, angle=45):
    if not bboxes:
        return image, bboxes
    # [bboxes] = bboxes
    # # print(bboxes)
    # bboxes = clip_bboxes(bboxes)
    # 将标签和边界框分开
    bboxes, labels = zip(*[(bbox[:-1], bbox[-1]) for bbox in bboxes])

    img_height, img_width = image.shape[:2]

    def clip_bbox(bboxes, img_width, img_height):
        clipped_bboxes = []
        for bbox in bboxes:
            # """将边界框修剪到图像范围内。"""
            x_min, y_min, x_max, y_max = bbox
            x_min = max(0, min(x_min, img_width))
            y_min = max(0, min(y_min, img_height))
            x_max = max(0, min(x_max, img_width))
            y_max = max(0, min(y_max, img_height))
            clipped_bboxes.append((x_min, y_min, x_max, y_max))
            # print(f'clipped_bboxes:{clipped_bboxes}')
            # print('--------------------')
        return clipped_bboxes

    bboxes = clip_bbox(bboxes, img_width, img_height)

    transform = A.Compose([
        A.Rotate(limit=(angle, angle), p=1)  # 旋转图像指定角度
    ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))


    # bboxes = clip_bboxes(bboxes)
    augmented = transform(image=image, bboxes=bboxes, labels=labels)

    # 将边界框和标签重新组合为一个列表
    augmented_bboxes = [bbox + (label,) for bbox, label in zip(augmented['bboxes'], augmented['labels'])]

    return augmented['image'], augmented_bboxes
def show_picinfo(image_path):
    image = cv2.imread(image_path)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    brightness = np.mean(gray)
    print('Brightness:', brightness)
    contrast = np.std(gray)
    print('Contrast:', contrast)

def fun_choice(index):
    if index == 0:
        return [False, False, False, False]
    elif index == 1:
        return [True, False, False, False]
    elif index == 2:
        return False, True, False, False
    elif index == 3:
        return False, False, True, False
    elif index == 4:
        return False, False, False, True
    elif index == 5:
        return True, True, False, False
    elif index == 6:
        return False, True, True, False
    elif index == 7:
        return True, False, True, False
    elif index == 8:
        return True, True, True, False
    elif index == 9:
        return True, True, True, True
    else:
        raise Exception("选择有误!")




def main(bright_change, noise, crop, rotat,brightness_value, contrast_value):
    input_folder = r"F:\desk\yolov5\data\dataset\images\val"
    image_output_folder = r"F:\desk\yolov5\data\dataset/newimages"

    bbox_folder = r"F:\desk\yolov5\data\dataset\labels\val"
    bbox_output_folder = r"F:\desk\yolov5\data\dataset/newlabels"

    if not os.path.exists(image_output_folder):
        os.makedirs(image_output_folder)
    if not os.path.exists(bbox_output_folder):
        os.makedirs(bbox_output_folder)
    for filename in os.listdir(input_folder):
        print(f'文件{filename}正在处理：')
        if filename.endswith(".jpg") or filename.endswith(".png") or filename.endswith(".bmp"):
            image_path = os.path.join(input_folder, filename)
            if filename.endswith(".jpg"):
                bbox_path = os.path.join(bbox_folder, filename.replace('.jpg', '.txt'))
            elif filename.endswith(".png"):
                bbox_path = os.path.join(bbox_folder, filename.replace('.png', '.txt'))
            elif filename.endswith(".bmp"):
                bbox_path = os.path.join(bbox_folder, filename.replace('.bmp', '.txt'))
            image = load_image(image_path)
            # 显示亮度对比度信息
            # show_picinfo(image_path)

            bboxes = load_bboxes(bbox_path, image.shape)
            # bboxes:[[528.99944, 653.000064, 867.0002400000001, 1067.000496, 0.0]]
            # print(f'bboxes:{bboxes}')
            # print('--------------------')

            # bboxes_with_labels = [(bbox[:-1], bbox[-1]) for bbox in bboxes]
            # bboxes_with_labels:[([528.99944, 653.000064, 867.0002400000001, 1067.000496], 0.0)]
            # print(f'bboxes_with_labels:{bboxes_with_labels}')
            # print('--------------------')


            if bright_change and not noise and not crop and not rotat:  # 改变亮度 yes
                print(f'---------任务：改变亮度---------')

                augmented_image, augmented_bboxes = change_brightness_contrast(image, bboxes,
                                                                               brightness_value, contrast_value)
                #   暗---->亮 brightness_value=0.01, contrast_value=2.2
                #   亮 ---->暗 brightness_value=-0.20, contrast_value=0.85
            elif noise and not bright_change and not crop and not rotat:  # 增加噪声 yes
                print(f'---------任务：增加噪声!---------')

                augmented_image, augmented_bboxes = augment_noise(image, bboxes)

            elif crop and not bright_change and not noise and not rotat:  # 裁剪 yes
                print(f'---------任务：裁剪!---------')

                augmented_image, augmented_bboxes = crop_image(image, bboxes, image.shape)

            elif rotat and not crop and not bright_change and not noise:  # 旋转 yes
                print(f'---------任务：旋转!---------')

                angle = random.choice([90, 180, 270, 360])
                augmented_image, augmented_bboxes = augment_rotation(image, bboxes, angle)

            elif bright_change and noise and not crop and not rotat:   # 亮度变化后加噪声 yes
                print(f'---------任务：亮度变化后加噪声!---------')
                #   暗---->亮 brightness_value=0.01, contrast_value=2.2
                #   亮 ---->暗 brightness_value=-0.20, contrast_value=0.85
                bright_image, bright_bboxes = change_brightness_contrast(image, bboxes,
                                                                         brightness_value, contrast_value)
                augmented_image, augmented_bboxes = augment_noise(bright_image, bright_bboxes)

            elif crop and noise and not bright_change and not rotat:  # 裁剪后加噪声 yes
                print(f'---------任务：裁剪后加噪声!---------')
                cropped_image, cropped_bboxes = crop_image(image, bboxes, image.shape)
                augmented_image, augmented_bboxes = augment_noise(cropped_image, cropped_bboxes)

            elif bright_change and crop and not noise and not rotat:  # 亮度变化后加裁剪 yes
                print(f'---------任务：亮度变化后加裁剪!---------')
                #   暗---->亮 brightness_value=0.01, contrast_value=2.2
                #   亮 ---->暗 brightness_value=-0.20, contrast_value=0.85
                # augmented_image, augmented_bboxes = augment_brightness_contrast(image, bboxes_with_labels, image.shape)
                bright_image, bright_bboxes = change_brightness_contrast(image, bboxes,
                                                                         brightness_value, contrast_value)
                augmented_image, augmented_bboxes = crop_image(bright_image, bright_bboxes, image.shape)

            elif bright_change and crop and noise and not rotat:  # 亮度变化后加裁剪加噪声 yes
                print(f'---------任务：亮度变化后加裁剪加噪声!---------')
                #   暗---->亮 brightness_value=0.01, contrast_value=2.2
                #   亮 ---->暗 brightness_value=-0.20, contrast_value=0.85
                bright_image, bright_bboxes = change_brightness_contrast(image, bboxes,
                                                                         brightness_value, contrast_value)
                # bright_bboxes = [(bbox[:-1], bbox[-1]) for bbox in bright_bboxes]
                a_image, a_bboxes = crop_image(bright_image, bright_bboxes, image.shape)
                augmented_image, augmented_bboxes = augment_noise(a_image, a_bboxes)
            elif rotat and bright_change and crop and noise:  # 旋转 亮度变化后 加裁剪 加噪声
                print(f'---------任务：旋转 亮度变化后 加裁剪 加噪声!---------')

                angle = random.choice([90, 180, 270, 360])
                rotated_image, rotated_bboxes = augment_rotation(image, bboxes, angle)
                bright_image, bright_bboxes = change_brightness_contrast(rotated_image, rotated_bboxes,
                                                                         brightness_value, contrast_value)
                a_image, a_bboxes = crop_image(bright_image, bright_bboxes, image.shape)
                augmented_image, augmented_bboxes = augment_noise(a_image, a_bboxes)
            else:
                raise Exception('参数设置错误!\n')

            save_image(augmented_image, image_output_folder, filename)
            augmented_bboxes = [list(bbox[:-1]) + [bbox[-1]] for bbox in augmented_bboxes]

            if filename.endswith(".jpg"):
                save_bboxes(augmented_bboxes, bbox_output_folder, filename.replace('.jpg', '.txt'),
                            augmented_image.shape)
            elif filename.endswith(".png"):
                save_bboxes(augmented_bboxes, bbox_output_folder, filename.replace('.png', '.txt'),
                            augmented_image.shape)
            elif filename.endswith(".bmp"):
                save_bboxes(augmented_bboxes, bbox_output_folder, filename.replace('.bmp', '.txt'),
                            augmented_image.shape)
            (f'文件{filename}处理完成')


if __name__ == "__main__":
    '''
    暗---->亮 brightness_value=0.01, contrast_value=2.2
    
    亮 ---->暗 brightness_value=-0.20, contrast_value=0.85
         0 -----------报错
         1 -----------任务：改变亮度---改亮度参数
         2 -----------任务：增加噪声
         3 -----------任务：裁剪
         4 -----------任务：旋转
         5 -----------任务：亮度变化后加噪声---改亮度参数
         6 -----------任务：裁剪后加噪声
         7 -----------任务：亮度变化后加裁剪----改亮度参数
         8 -----------任务：亮度变化后加裁剪加噪声----改亮度参数
         9 -----------任务：旋转 亮度变化后 加裁剪 加噪声---改亮度参数
    '''
    bright_change, noise, crop, rotat = fun_choice(9)
    main(bright_change, noise, crop, rotat, brightness_value=-0.10, contrast_value=0.85)
